/*
 *  Copyright (C) 2010 Matthias Buch-Kromann <mbk.isv@cbs.dk>
 * 
 *  This file is part of the MatrixParser package.
 *  
 *  The MatrixParser program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Lesser General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 * 
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Lesser General Public License for more details.
 * 
 *  You should have received a copy of the GNU Lesser General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package org.osdtsystem.matrixparser.parsers;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import org.osdtsystem.matrixparser.data.CONLLSentence;
import org.osdtsystem.matrixparser.featureextraction.UnionExtractor;
import org.osdtsystem.matrixparser.parsers.output.DependencyTree;
import org.osdtsystem.matrixparser.featureextraction.HigherOrderExtractor;
import org.osdtsystem.matrixparser.featureextraction.FirstOrderExtractorLabelled;
import org.osdtsystem.matrixparser.featureextraction.FirstOrderExtractorUnlabelled;
import org.osdtsystem.matrixparser.featureextraction.HigherOrderMbkCframes;
import org.osdtsystem.matrixparser.features.FeatureVector;
import org.osdtsystem.matrixparser.features.FeatureVectorFactory;
import org.osdtsystem.matrixparser.features.MstCache;
import org.osdtsystem.matrixparser.features.TwoWayFeatureHandler;
import org.osdtsystem.matrixparser.learners.Scorer;
import org.osdtsystem.matrixparser.features.WeightVector;
import org.osdtsystem.matrixparser.logging.Log;
import org.osdtsystem.matrixparser.main.Options;

/**
 *
 * @author Matthias Buch-Kromann <mbk.isv@cbs.dk>
 */
public class IncrementalLabMSTParser extends LabelledMSTParser {
    // Extractors
    HigherOrderExtractor higherOrderExtractor;
    
    // Parameters
    int klabels;
    int kedges;    //
    int kheads;      // The number of different heads each node may try for TryAttachRight
    int[] allEdgeIndices;

//    DependencyTree tree = null;
//    private WeightVector weights;
    //int lastnode = 0;
    //int leftscan = 0;
    
    public IncrementalLabMSTParser(String name, Scorer scorer, HigherOrderExtractor higherOrderExtractor,
            FirstOrderExtractorUnlabelled unlabelledMstExtractor,
            FirstOrderExtractorLabelled labelledMstExtractor,
            TwoWayFeatureHandler labelAlphabet) {
        super(name, scorer, unlabelledMstExtractor, labelledMstExtractor, labelAlphabet);
        this.higherOrderExtractor = higherOrderExtractor;

        // Read ktypes and kedges
        klabels = Options.intOption("incparse.ktypes");
        kedges = Options.intOption("incparse.kedges");
        kheads = Options.intOption("incparse.kheads");

        allEdgeIndices = intseq(0, klabels);

        Log.config("IncLabMstParser: ktypes=" + klabels + " kedges=" + kedges + " kheads=" + kheads);
    }

    public IncrementalLabMSTParser(String name, Scorer scorer, ParsingModel model, UnionExtractor extractor) {
        this(name, scorer, extractor, extractor, extractor, model.labelHandler());
    }

    @Override
    public DependencyTree parse(CONLLSentence sentence, WeightVector weightVector) {
        // Create dependency tree for sentence, parameters, and cache
        DependencyTree tree = new DependencyTree(sentence, labelAlphabet, lastComplementId);
        ParseState state = new ParseState(tree, weightVector);

        // Compute edge-factored labelled and unlabelled scores
        ActionList<Action> actions = new ActionList<Action>(state.cklabels);
        double[] headscores = new double[state.nodes];
        for (int dependent = 0; dependent < state.nodes; dependent++) {
            for (int head = 0; head < state.nodes; head++) {
                double unlabScore = scorer.getScore(weightVector,
                        state.cache.unlabelledFeatures(head, dependent));

                // Find k-best labels
                actions.clear();
                for (int t = 0; t < labelAlphabet.alphabetSize(); t++) {
                    FeatureVector headFV = state.cache.labelledHeadFeatures(head, dependent, t);
                    FeatureVector depFV = state.cache.labelledDependentFeatures(head, dependent, t);
                    double labScore = scorer.getScore(weightVector, headFV) + scorer.getScore(weightVector, depFV);
                    double score = unlabScore + labScore;
                    if (score > actions.worstScore())
                        actions.add(new Action(score, t));
                }

                // Store k-best labels
                int imax = Math.min(state.cklabels, actions.size());
                for (int i = 0; i < imax; ++i) {
                    Action action = actions.get(i);
                    int arrayPos = state.index(head, dependent, i);
                    state.scoresK[arrayPos] = action.score();
                    state.labelsK[arrayPos] = action.get(0);
                }

                headscores[head] = state.scoresK[state.index(head, dependent, 0)];
            }

            // Store k-best score for head
            Arrays.sort(headscores);
            state.kbest1stOrderScore[dependent] = headscores[Math.max(0, headscores.length - kheads)];
        }


        // Create root node
        state.tree.addHistory("AddRoot", new ActionAddNode(0, sentence));

        // Process nodes incrementally from left to right
        for (int node = 1; node < state.nodes; node++) {
            // Set current position
            tree.setLastnode(node);
            tree.setLeftscan(node);

            // Rescore preceding node cframes as they mature
            for (int offset : HigherOrderMbkCframes.offsets) {
                if (node - offset >= 0)
                    updateSecondOrderScores(state, node - offset);
            }

            // Attach latest node to the best possible left head
            attachToBestLeft(state, node);

            // Now try to see if any of the nodes to the left want to attach to the right
            for (int dependent = node - 1; dependent > 0; --dependent) {
                tree.setLeftscan(dependent);
                if (state.scores1stOrder[dependent] < state.kbest1stOrderScore[dependent]) {
                    // Attachment to new node would be better: try to repair the
                    // old attachment
                    tryAttachToRight(state, dependent, node);
                }
            }
        }

        // Create output
        return tree;
    }

    void attachToBestLeft(ParseState state, int node) {
        // Create action lists
        ActionList<Action> efActions = new ActionList<Action>(kedges);

        // Find best attachment actions wrt. edge-factored scores
        for (int head = node - 1; head >= 0; --head) {
            for (int i = 0; i < state.cklabels; ++i) {
                double score = state.scoresK[state.index(head, node, i)];
                if (score > efActions.worstScore()) {
                    efActions.add(new Action(score, head, i));
                } else {
                    break;
                }
            }
        }

        // Compute totally best action wrt. total scores, among all surviving actions
        ActionList<ActionMoveEdge> bestAction = new ActionList<ActionMoveEdge>(1);
        for (int i = 0; i < efActions.size(); ++ i) {
            // Retrieve action
            Action action = efActions.get(i);
            int head = action.get(0);
            int[] typelist = {action.get(1)};
        
            // Add move-edge action to 1-best list
            addMoveEdgeActionsToActionList(state, node, head, typelist, bestAction);
        }

        // Effectuate the best action if it improves the score
        ActionMoveEdge action = bestAction.get(0);
        action.log(" [L]");
        action.apply(this, state);
        state.tree.addHistory("AttachLeft", new ActionAddNode(node, state.tree.sentence()), action);
    }

    void tryAttachToRight(ParseState state, int node, int head) {
        if (node == head) {
            return;
        } else if (! state.tree.dominates(node, head)) {
            // No problems with cycles: just try to reattach directly
            tryAttachToRightNoCycle(state, node, head);
        } else {
            // Oops! node dominates head so direct attachment would result in a cycle.
            // Try to see if we can reattach any of the nodes from head and up to
            // (and excluding) node to a head that isn't dominated by node,
            // at a lower cost than the potential gain from reattaching node
            // to head
            tryAttachToRightCycle(state, node, head);
        }
    }

    void tryAttachToRightNoCycle(ParseState state, int node, int head) {
        // Create action list and find best action
        ActionList<ActionMoveEdge> bestAction = new ActionList<ActionMoveEdge>(1);
        addMoveEdgeActionsToActionList(state, node, head, allEdgeIndices, bestAction);

        // Effectuate the best action if it improves the score
        ActionMoveEdge action = bestAction.get(0);
        if (action.score() > 0) {
            action.log(" [RNC]");
            action.apply(this, state);
            state.tree.addHistory("RepairRightNoCycle", action);
        }
    }

    void tryAttachToRightCycle(ParseState state, int node, int head) {
        // Find best node action, and abandon repair immediately if score is negative
        // (presumably, all nodes on the path have their optimal attachment already)
        ActionList<ActionMoveEdge> bestNodeAction = new ActionList<ActionMoveEdge>(1);
        addMoveEdgeActionsToActionList(state, node, head, allEdgeIndices, bestNodeAction);
        ActionMoveEdge nodeAction = bestNodeAction.get(0);
        double deltaNodeMove = nodeAction.score();
        if (deltaNodeMove < 0)
            return;

        // Node dominates head in the old graph. So find the best reattachment actions
        // (wrt. edge-factored scores) for nodes on the path from head (inclusive) up 
        // to node (exclusive) where the new head is not dominated by node in the old graph.
        ActionList<Action> efPathActions = new ActionList<Action>(kedges);
        int pnode = head, phead = 0;
        while (pnode != node && pnode != 0) {
            // Calculate pnode ef-score
            double efScorePnodeOld = state.scores1stOrder[pnode];

            // Calculate the price of reattaching repaired somewhere that
            // isn't dominated by node
            for (phead = state.tree.lastnode; phead >= 0; --phead) {
                if (phead != node && ! state.tree.dominates(node, phead)) {
                    for (int i = 0; i < state.cklabels; ++i) {
                        double efScorePnodeNew = state.scoresK[state.index(phead, pnode, i)];
                        double deltaEF = efScorePnodeNew - efScorePnodeOld;
                        if (deltaEF > efPathActions.worstScore())
                            efPathActions.add(new Action(deltaEF, pnode, phead, i));
                    }
                }
            }

            // Now evaluate the best edge-factored Move up to the next node on the path from head to node
            pnode = state.tree.head(pnode);
        }

        // Apply the best node action
        nodeAction.apply(this, state);

        // Find the best path action wrt. total scores, among the surviving path actions
        ActionList<ActionMoveEdge> bestPathAction = new ActionList<ActionMoveEdge>(1);
        for (int i = 0; i < efPathActions.size(); ++i) {
            Action action = efPathActions.get(i);
            pnode = action.get(0);
            phead = action.get(1);
            int[] pindex = {action.get(2)};
            addMoveEdgeActionsToActionList(state, pnode, phead, pindex, bestPathAction);
        }

        // Apply best path action if score improves, otherwise restore original
        // graph by unapplying the best node action
        ActionMoveEdge pathAction = bestPathAction.get(0);
        double deltaPnodeMove = pathAction.score();
        if (deltaNodeMove + deltaPnodeMove > 0) {
            nodeAction.log(" [RC1]");
            pathAction.log(" [RC2]");
            pathAction.apply(this, state);
            state.tree.addHistory("RepairRightCycle", nodeAction, pathAction);
        } else {
            nodeAction.unapply(this, state);
        }
    }

    void addMoveEdgeActionsToActionList(ParseState state, int node, int newhead, int[] indices,
            ActionList<ActionMoveEdge> actionList) {
        // Backup old configuration
        DependencyTree tree = state.tree;
        int oldhead = tree.head(node);
        int oldtype = tree.label(node);
        double oldEFscore = state.scores1stOrder[node];
        //System.err.println(" addMoveEdge: node=" + node + " nhead=" + newhead
        //        + " ohead=" + oldhead + " otype=" + oldtype + " ind=" + toString(indices));

        // Compute scores for given edge indices
        for (int i : indices) {
            // Retrieve new type and edgefactored score
            int index = state.index(newhead, node, i);
            int newtype = state.labelsK[index];
            double newEFscore = state.scoresK[index];
            double deltaScore = newEFscore - oldEFscore;

            // Update dependency tree and calculate affected nodes (node, new and old head/siblings)
            int[] affectedCustom = {node, oldhead, newhead};
            int[] affectedOld = affectedSiblings(state, node, affectedCustom);
            tree.setDependency(node, newhead, newtype);
            int[] affected = affectedSiblings(state, node, affectedOld);

            // Update scores for affected nodes
            double[] oldScores = saveSecondOrderScores(state, affected);
            updateSecondOrderScores(state, affected);
            double[] newScores = saveSecondOrderScores(state, affected);

            // Update total score with nonlocal scores
            for (int j = 0; j < newScores.length; ++j)
                deltaScore += newScores[j] - oldScores[j];

            // Restore dependency tree
            if (oldhead >= 0) {
                tree.setDependency(node, oldhead, oldtype);
            } else {
                tree.removeDependency(node, newhead, newtype);
            }
            restoreNonEFScores(state, affected, oldScores);

            // Update action list
            if (deltaScore > actionList.worstScore()) {
                actionList.add(new ActionMoveEdge(deltaScore, oldEFscore, newEFscore,
                        node, oldhead, oldtype, newhead, newtype,
                        affected, newScores, oldScores));
            }
        }
    }

    String toString(int[] list) {
        StringBuilder sb = new StringBuilder("[");
        String sep = "";
        for (int i : list)
            sb.append(sep).append(i);
        sb.append("]");
        return sb.toString();
    }

    int[] affectedSiblings(ParseState state, int node, int... preaffected) {
        // Find head and return affected if empty
        DependencyTree tree  = state.tree;
        int head = tree.head(node);
        if (head < 0)
            return preaffected;

        // Create list of affected nodes
        List<Integer> affected = new ArrayList<Integer>();
        for (int i : preaffected)
            affected.add(i);

        // Find label
        int label = tree.label(node);

        // Find position of node in children
        List<Integer> children = node < head ? tree.leftChildren(head) : tree.rightChildren(head);
        int inode = 0;
        for (inode = 0; children.get(inode) != node && inode < children.size(); ++inode) { }

        // Find affected siblings
        if (! tree.isComplementRelation(label)) {
            // Adjuncts only affect adjacent siblings
            if (inode > 0)
                affected.add(children.get(inode - 1));
            if (inode < children.size() - 1)
                affected.add(children.get(inode + 1));
        } else {
            // Complements affect all siblings until the nearest complement siblings
            // Left side
            for (int i = inode - 1; i >= 0; --i) {
                int sibling = children.get(i);
                affected.add(sibling);
                if (tree.isComplement(sibling))
                    break;
            }

            // Right side
            for (int i = inode + 1; i < children.size(); ++i) {
                int sibling = children.get(i);
                affected.add(sibling);
                if (tree.isComplement(sibling))
                    break;
            }
        }

        // Create array
        int[] result = new int[affected.size()];
        for (int i = 0; i < result.length; ++i)
            result[i] = affected.get(i);
        return result;
    }

    FeatureVector[] updateSecondOrderScores(ParseState state, int... nodes) {
        FeatureVector[] fvs = new FeatureVector[nodes.length];
        for (int i = 0; i < nodes.length; ++i) {
            int node = nodes[i];
            fvs[i] = null;
            higherOrderExtractor.extractHigherOrder(fvs[i], state.tree, node);
            if (node >= 0)
                state.scores2ndOrder[node] = scorer.getScore(state.weights, fvs[i]);
        }
        return fvs;
    }

    double[] saveSecondOrderScores(ParseState state, int[] nodes) {
        double[] scores = new double[nodes.length];
        for (int i = 0; i < nodes.length; ++i) {
            int node = nodes[i];
            scores[i] = node >= 0 ? state.scores2ndOrder[node] : 0;
        }
        return scores;
    }

    void restoreNonEFScores(ParseState state, int[] nodes, double[] scores) {
        for (int i = 0; i < nodes.length; ++i) {
            int node = nodes[i];
            if (node >= 0) {
                state.scores2ndOrder[node] = scores[i];
            }
        }
    }

    final int[] intseq(int start, int stop) {
        int[] result = new int[stop-start];
        int val = start;
        for (int i = result.length - 1; i >= 0; --i)
            result[i] = start + val;
        return result;
    }

    @Override
    public List<FeatureVector> featureVectors(DependencyTree tree,
            List<FeatureVector> list) {
        // Add MST feature vectors
        super.featureVectors(tree, list);

        // Add non-edgefactored feature vectors
        for (int node = 0; node < tree.nodes(); ++node) {
            FeatureVector fv = FeatureVectorFactory.makeNodeFeatureVector();
            higherOrderExtractor.extractHigherOrder(fv, tree, node);
            list.add(fv);
        }

        // Return
        return list;
    }

    // Stored graph information for parse
    class ParseState {
        final WeightVector weights;
        final DependencyTree tree;
        final int nodes;
        final int cklabels;
        final double[] scoresK;
        final int[] labelsK;
        final double[] scores1stOrder;
        final double[] scores2ndOrder;
        final double[] kbest1stOrderScore;
        MstCache cache;

        ParseState(DependencyTree tree, WeightVector weights) {
            this.tree = tree;
            this.weights = weights;
            this.nodes = tree.nodes();
            cklabels = Math.min(klabels, labelAlphabet.alphabetSize());
            scoresK = new double[nodes * nodes * cklabels];
            labelsK = new int[nodes * nodes * cklabels];
            scores1stOrder = new double[nodes];
            scores2ndOrder = new double[nodes];
            kbest1stOrderScore = new double[nodes];
            Arrays.fill(scores1stOrder, 0);
            Arrays.fill(scores2ndOrder, 0);
            cache = mstCache(tree.sentence());
        }
        
        final int index(int head, int dependent, int edgeindex) {
            if (head < 0 || head >= nodes)
                throw new IllegalArgumentException("Illegal head: " + head + " (should be 0..." + (nodes-1) + ")");
            if (dependent < 0 || dependent >= nodes)
                throw new IllegalArgumentException("Illegal dependent: " + dependent + " (should be 0..." + (nodes-1) + ")");
            if (edgeindex < 0 || edgeindex >= cklabels)
                throw new IllegalArgumentException("Illegal edge index: " + edgeindex + " (should be 0..." + (cklabels-1) + ")");
            return (head * nodes + dependent) * cklabels + edgeindex;
        }
    }

}